import json

from openai import OpenAI
from python_execute import PythonExecute
from terminate import Terminate
from file_saver import FileSaver
from tool_collection import ToolCollection

client = OpenAI(base_url="http://192.168.2.112:8080/v1", api_key="helloworld")

tool_collection = ToolCollection(FileSaver(), PythonExecute(), Terminate())
print("all tools:", tool_collection.to_params())

while True:
    messages = [{"role": "user", "content": ""}]
    messages[0]["content"] = input("please input anything prompt")
    if messages[0]["content"].strip() == "exit":
        break
    completion = client.chat.completions.create(
        model="gpt-4o",
        messages=messages,
        tools=tool_collection.to_params(),
        tool_choice="required"
    )
    print(completion)
    agent_chooses = completion.choices
    agent_chose_count = len(agent_chooses)
    for chose in agent_chooses:
        agent_chose_message = chose.message
        print(agent_chose_message)
        print(agent_chose_message.content)
        agent_tool_calls = agent_chose_message.tool_calls
        print("agent call counts of tools:", len(agent_tool_calls))
        for tool_call in agent_tool_calls:
            call_function = tool_call.function
            call_function_name = call_function.name
            call_function_arguments = call_function.arguments
            print(call_function_name, call_function_arguments)
            function_in_tool_collection = tool_collection.tool_map[call_function_name]
            print(function_in_tool_collection)
            if function_in_tool_collection is None:
                print("no function in tool collection")
            else:
                arguments = json.loads(call_function_arguments or "{}")
                ret = function_in_tool_collection.execute(**arguments)
                print(ret)
